from app.model import Model
import app.chunking as chunking
import chromadb
import os


class Database:
    QUERY_RESULT = 5
    DATABASE_PATH = "./data/chroma.db"

    def __init__(self, model: Model, create: bool=True):
        self.db_client = chromadb.PersistentClient(Database.DATABASE_PATH)
        self.db_collection = self.db_client.get_or_create_collection("linghuchong")
        self.model = model

        if create:
            self.create()

    def create(self) -> None:
        for index, chunk in enumerate(chunking.get_chunks()):
            print(f'Process: {chunk}')
            embedding = self._embed(chunk)
            self.db_collection.upsert(
                ids=str(index),
                documents=chunk,
                embeddings=embedding
            )

    def query(self, question: str) -> list[str]:
        question_embedding = self._embed(question)
        result = self.db_collection.query(
            query_embeddings=question_embedding,
            n_results=Database.QUERY_RESULT
        )

        return result["documents"][0]

    def _embed(self, text: str) -> list[str]:
        return self.model.embed(text)